import torch
import torch.nn as nn
import torch.nn.functional as F

class TotalVariation(nn.Module):
    def __init__(self,):
        super(TotalVariation,self).__init__()

    def forward(self, comps, mask):
        comps_hole= comps * (1-mask)
        comps_hole_iplus1 = comps_hole[:, :, 1:, :]
        comps_hole_jplus1 = comps_hole[:, :, :, 1:]
        
        comps_hole_iplus1 = torch.cat((comps_hole_iplus1, torch.zeros(comps.size(0), comps.size(1), 1, comps.size(3))), dim=2)
        comps_hole_jplus1 = torch.cat((comps_hole_jplus1, torch.zeros(comps.size(0), comps.size(1), comps.size(2), 1)), dim=3)

        loss = F.l1_loss(comps_hole_iplus1, comps, reduction='sum') + F.l1_loss(comps_hole_jplus1, comps, reduction='sum')

        return loss


